# coding: utf-8

import logging
import re
import string
import xml.dom.minidom

log = logging.getLogger('dsMatX_ResolveEditor')
from Katana import NodegraphAPI, DrawingModule, Utils

import domtree as dom

class BColors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

    def disable(self):
        self.HEADER = ''
        self.OKBLUE = ''
        self.OKGREEN = ''
        self.WARNING = ''
        self.FAIL = ''
        self.ENDC = ''


class BaseNode(object):

    def __init__(self, name, type=None, parent=None):

        super(BaseNode, self).__init__()

        #These is one "name" argument at least,

        # Get the node from nodegraph
        self._node = NodegraphAPI.GetNode(name)

        if self._node is None:

            #The second argument is the type of new node to create,
            #the third argument is the parent of the node
            with NodegraphAPI.InteractiveContext():

                self._node = NodegraphAPI.CreateNode(type, parent)
                # self._node._flushAll()


            self.name = name

    @property
    def color(self):

        return DrawingModule.SetCustomNodeColor(self.node)

    @color.setter
    def color(self, color):

        DrawingModule.SetCustomNodeColor(self.node, color[0], color[1], color[2])

    @property
    def node(self):

        return self._node

    @node.setter
    def node(self, n):

        self._node = n

    @property
    def name(self):

        try:
            return self._node.getParameter('name').getValue(0)

        except:

            self._node.getName(0)

    @name.setter
    def name(self, name):

        # Try to set the parameter "name" first, then to set node name.
        # In katana, the node name will be update as the parameter "name"
        try:

            self._node.getParameter('name').setValue(str(name), 0)

        except:

            self._node.setName(str(name))


    def addInputPort(self, name):
        #Add an input port using the name
        return self._node.addInputPort(name)

    def outputPort(self, name):
        #Get output port by the name
        return self._node.getOutputPort(name)

    def inputPort(self, name):
        return self._node.getInputPort(name)

    @property
    def numInputPorts(self):
        #Get the numbers of the input ports
        return self._node.getNumInputPorts()


    def inputPortByIndex(self, id=0):
        #Get the input port by id
        return self._node.getInputPortByIndex(id)


    def outputPortByIndex(self, id=0):
        #Get the output port by id
        return self._node.getOutputPortByIndex(id)


class KatanaNode(BaseNode):
    def __init__(self, name, type, parent):
        super(KatanaNode, self).__init__(name, type, parent)

    def addShaderInputPort(self, rendererName, shaderType):
        self._node.addShaderInputPort(rendererName, shaderType)


class MaterailAssignNode(KatanaNode):

    def __init__(self, name, type, parent):
        super(MaterailAssignNode, self).__init__(name, type, parent)

    @property
    def expression(self):
        return self._node.getParameter('args.materialAssign.value').getValue(0)

    @expression.setter
    def expression(self, exp):
        self._node.getParameter('args.materialAssign.value').setExpression("scenegraphLocationFromNode(getNode('%s'))" % exp)

    @property
    def CEL(self):
        return self._node.getParameter('CEL').getValue(0)

    @CEL.setter
    def CEL(self, cel):
        self._node.getParameter('CEL').setValue(cel, 0)


class ArnoldNode(BaseNode):

    nodeType_map = {
        'blendColors': 'mix_rgba',
        'file': 'image',
        'ramp': 'ramp_rgb',
        'place2dTexture': 'uv_transform',
        'aiStandardSurface': 'standard_surface',
        'checker': 'checkerboard',

        'aiBump2d': 'bump2d',
        'aiBump3d': 'bump3d',
        'aiNormalMap': 'normal_map',
        'aiRange': 'range',
        'aiFacingRatio': 'facing_ratio',
        'aiSwitch': 'switch_shader',
        'aiColorToFloat': 'rgb_to_float',
        'aiColorCorrect': 'color_correct',
    }

    param_map = {
        'specularIOR': 'specular_IOR',
        'coatIOR': 'coat_IOR',
        'thinFilmIOR': 'thin_film_IOR',
        'repeatUV': 'repeat',
        'rotateUV': 'rotate',
        'noiseUV': 'noise',
        'offsetUV': 'offset',
        'bumpDepth': 'bump_height',
        'colors': 'color',
        'colorPositions': 'position',
        'colorSpace': 'color_space',
        'normalCamera': 'normal',
        'outValue': 'out',
        'output': 'out',
        'outNormal': 'out',
        'outColor': 'out',
        'outColorR': 'out.r',
        'outColorG': 'out.g',
        'outColorB': 'out.b',
        'outAlpha': 'out',
        'blender': 'mix',
        'subsurfaceType': 'subsurface_type',
        'aiFlipR': 'invert_x',
        'aiFlipG': 'invert_y',
        'aiFlipB': 'invert_z',
        # 'colorGain': 'multiply',
        # 'colorOffset': 'add',
        # 'alphaGain': 'alpha_multiply',
        # 'alphaOffset': 'alpha_add',

        'color1': ['input1', 'color1'],
        'color2': ['input2', 'color2'],
        'bumpValue': ['input', 'bump_map'],
    }
    ramp_map = {
        'type': {
            "0": "v",
            "1": "u",
            "2": "diagonal",
            "3": "radial",
            "4": "circular",
            "5": "box",
            "6": "custom",
            "7": "custom",
            "8": "custom",
        },
        'interpolation': {
            '0': '0',
            '1': '1',
            '2': '3',
            '3': '3',
            '4': '2',
            '5': '1',
            '6': '1'
        }
    }
    def __init__(self, name, type=None, parent=None):

        if type and parent:
            super(ArnoldNode, self).__init__(name, 'ArnoldShadingNode', parent)
            self.nodeType = type
        else:
            super(ArnoldNode, self).__init__(name)

    @property
    def nodeType(self):
        return self._node.getParameter('nodeType').getValue(0)

    @nodeType.setter
    def nodeType(self, type):
        if ArnoldNode.nodeType_map.has_key(type):
            type = ArnoldNode.nodeType_map[type]
        if type.startswith('ai'):
            type = type[2:].lower()

        with NodegraphAPI.InteractiveContext():
            self._node.getParameter('nodeType').setValue(str(type), 0)
            self._node.checkDynamicParameters()



    def setValue(self, params_xml):
        params_xml = self.verifyParamsXML(params_xml)

        for param in params_xml:
            #Get the real parameter name
            parame_name = param.getAttribute('name')[2:]

            if parame_name[-1] in ['R', 'G', 'B', 'X', 'Y', 'Z'] and not parame_name.endswith('IOR'):
                # If a component parameter, then continue. Because we will set enable through it's parent parameter
                continue

            else:
                if self._node.getParameter('parameters.%s' %ReformatParameter(parame_name, False)):
                    parame_name = ReformatParameter(parame_name, False)

                elif ArnoldNode.param_map.has_key(parame_name):
                    parame_name = ArnoldNode.param_map[parame_name]

                else:
                    print (BColors.HEADER + 'Discard to set parameter [%s] for node [%s]' %(parame_name, self.name) + BColors.ENDC)
                    continue

            print (BColors.OKBLUE + 'Setting parameter [%s] for node [%s]' % (parame_name, self.name) + BColors.ENDC)

            data_type = param.getAttribute('type')
            data_value = str(param.getAttribute('value'))
            param_knode = self._node.getParameter('parameters.%s' % parame_name)

            #----------------------------------------------------------------------------------------------------------
            # print (BColors.HEADER + '%s,\t %s,\t %s,\t %s' %(data_type, data_value, parame_name, self.name) + BColors.ENDC)
            # ----------------------------------------------------------------------------------------------------------

            param_knode.getChild('enable').setValue(True, 0)

            #String
            if data_type == 'string':

                if self.nodeType == 'image':

                    if parame_name == 'color_space':
                        param_knode.getChild('value').setValue('linear', 0)
            #Integer
            elif data_type == 'integer':

                if self.nodeType.startswith('ramp'):
                    if parame_name == 'type':
                        # Only for ramp  of node
                        param_knode.getChild('value').setValue(ArnoldNode.ramp_map['type'][data_value], 0)

                elif self.nodeType == 'standard_surface':
                    # The type of "subsurface_type" is string
                    param_knode.getChild('value').setValue(str(data_value), 0)

            elif data_type == 'boolean':

                if param_knode.getChild('value') == 'False':
                    # Because True is 1, and False is 0 in Katana bool value
                    param_knode.getChild('value').setValue(0, 0)
                else:
                    param_knode.getChild('value').setValue(1, 0)

            elif data_type == 'Color3ArrayAttribute':
                continue

            elif data_type == 'FloatArrayAttribute':

                if self.nodeType.startswith('ramp'):

                    if parame_name == 'position':
                        rampColorList = list()
                        positionsList = list()
                        colorByPosition = dict()
                        indexInParamList = getIndexInList(params_xml, 'name', 'i_interpolation')
                        interp = 2

                        if indexInParamList:
                            # Maybe parameter "interpolation" does not exits
                            interp = ArnoldNode.ramp_map['interpolation'][params_xml[indexInParamList].getAttribute('value')]
                        indexInParamList = getIndexInList(params_xml, 'name', 'i_colors')
                        i_colors = params_xml[indexInParamList].getAttribute('value').split(' ')

                        for color in i_colors:
                            color = color.split(',')
                            rampColorList.append(tuple((float(f) for f in color)))
                        i_colorPositions = str(param.getAttribute('value')).split(',')
                        positionsList = [float(f) for f in i_colorPositions]

                        for i in range(len(positionsList)):
                            colorByPosition[positionsList[i]] = rampColorList[i]

                        # print colorByPosition
                        with NodegraphAPI.InteractiveContext():
                            rampParameter = param_knode.getParent().getParent().getChild('parameters')
                            positions = sorted(colorByPosition)
                            n = len(positions)
                            Utils.UndoStack.OpenGroup('Set %d Colors on %s' % (n, rampParameter.getFullName()))

                            try:
                                ramp = rampParameter.getChild('ramp')
                                ramp.getChild('enable').setValue(1, 0)
                                rampValue = ramp.getChild('value')
                                rampValue.setValue(n + 2, 0)

                                # interValue =
                                interpParameter = rampParameter.getChild('interpolation')
                                interpParameter.getChild('enable').setValue(1, 0)
                                interpValueParameter = interpParameter.getChild('value')
                                interpValueParameter.resizeArray(n + 2)
                                interpValueParameter.getChildByIndex(0).setValue(interp, 0)

                                positonParameter = rampParameter.getChild('position')
                                positonParameter.getChild('enable').setValue(1, 0)
                                positonValueParameter = positonParameter.getChild('value')
                                positonValueParameter.resizeArray(n + 2)
                                positonValueParameter.getChildByIndex(0).setValue(positions[0], 0)

                                for i, position in enumerate(positions):
                                    interpValueParameter.getChildByIndex(1 + i).setValue(interp, 0)
                                    positonValueParameter.getChildByIndex(1 + i).setValue(position, 0)
                                # #
                                interpValueParameter.getChildByIndex(n + 1).setValue(interp, 0)
                                positonValueParameter.getChildByIndex(n + 1).setValue(positions[-1], 0)

                                # colorsParameterName = '%s_Colors' % parameterName
                                colorsParameter = rampParameter.getChild('color')
                                colorsParameter.getChild('enable').setValue(1, 0)
                                colorsValueParameter = colorsParameter.getChild('value')
                                colorsValueParameter.resizeArray(3 * (n + 2))
                                for i, position in enumerate(positions):
                                    r, g, b = colorByPosition[position]
                                    colorsValueParameter.getChildByIndex(3 + 3 * i).setValue(r, 0)
                                    colorsValueParameter.getChildByIndex(3 + 3 * i + 1).setValue(g, 0)
                                    colorsValueParameter.getChildByIndex(3 + 3 * i + 2).setValue(b, 0)
                                r, g, b = colorByPosition[positions[0]]
                                colorsValueParameter.getChildByIndex(0).setValue(r, 0)
                                colorsValueParameter.getChildByIndex(1).setValue(g, 0)
                                colorsValueParameter.getChildByIndex(2).setValue(b, 0)
                                r, g, b = colorByPosition[positions[-1]]
                                colorsValueParameter.getChildByIndex(3 * (n + 1)).setValue(r, 0)
                                colorsValueParameter.getChildByIndex(3 * (n + 1) + 1).setValue(g, 0)
                                colorsValueParameter.getChildByIndex(3 * (n + 1) + 2).setValue(b, 0)

                            finally:
                                Utils.UndoStack.CloseGroup()

            # Color3, Float3, Vector2
            elif data_type[-1].isdigit():
                params = []

                for i in range(int(data_type[-1])):
                    params.append(param_knode.getChild('value').getChild('i%d' % i))

                if None not in params:
                    v = param.getAttribute('value').split(',')

                    for index, p in enumerate(params):
                        p.setValue(v[index], 0)
            else:
                param_knode.getChild('value').setValue(data_value, 0)


    def verifyParamsXML(self, params_xml):
        #Parameter: params_xml, the parameters elements of Current node.

        if self.nodeType == 'bump2d':
            attrs_names = [param.getAttribute('name') for param in params_xml]

            if 'i_bumpInterp' in attrs_names:
                indexParamInList = getIndexInList(params_xml, 'name', 'i_bumpInterp')

                if params_xml[indexParamInList].getAttribute('value') != 0:
                    self.nodeType = 'normal_map'
                    #----------------------------------------------------------------------------------
                    #Parameter name need to change, and must be modified.
                    index = getIndexInList(params_xml, 'name', 'i_bumpDepth')
                    if index != None:
                        params_xml[index].setAttribute('name', 'i_strength')
                    # ----------------------------------------------------------------------------------
                params_xml.pop(indexParamInList)

        return params_xml


def ResolveMatXS(gnode, filePath, location, matLocation):
    '''
    @param gnode:               supertools.node
    @param filePath:            the path to XML file
    @param location:            the location of Assets
    @param matLocation:         the location of XML file
    @return:
    '''
    print '\n--------------------------------- dsMatX_Resolve beginning -------------------------------------------\n'

    log.info('Reading ' + filePath)

    xml_dom  = xml.dom.minidom.parse(filePath)
    mtlXS_xml = prefix_gainOffset(xml_dom)
    mtlXS_xml = dom.main(mtlXS_xml)
    # mtlXS_xml = prefix_place2dTexture(mtlXS_xml)


    if len(mtlXS_xml.getElementsByTagName('materialx')) == 0:
        log.warning('File is not a valid MaterialX file.')
        return

    MAT_NAME_SPACE = matLocation
    #Create the main merge node
    mtlXmerge = KatanaNode(name=gnode.getName() + '_Merge', type='Merge', parent=gnode)
    mtlXmerge.color = [0.569, 0.449, 0.089]
    #Get all shading group in the xml file, Material = Maya Shading Group
    shadingGroup_xmlNodes = mtlXS_xml.getElementsByTagName('material')
    #Get all shaders in the xml file, Shader =  Maya Shader Node
    shaderNameToType = dict()
    shader_xmlNodes = mtlXS_xml.getElementsByTagName('shader')

    for node in shader_xmlNodes:
        shader_name = node.getAttribute('name')
        shader_type = node.getAttribute('shaderprogram')
        shaderNameToType[shader_name] = shader_type

    materialAssign_kNodes = list()
    shadingGroup_kNodes = list()

    for x, sg in enumerate(shadingGroup_xmlNodes):
        shadingGroup_name = MAT_NAME_SPACE + '_' + str(sg.getAttribute('name'))
        shadingGroup = KatanaNode(name=shadingGroup_name + '_GRP', type='Group', parent=gnode)
        shadingGroup.color = [0.17, (x + 0.2) / (len(shadingGroup_xmlNodes) * 2), 0.29]
        #A displacement node is exist if the number of 'shaderref' tag is more than one
        hasDisp = len(sg.getElementsByTagName('shaderref')) > 1
        network_knode = KatanaNode(name=shadingGroup_name, type='NetworkMaterial', parent=shadingGroup.node)
        network_knode.addShaderInputPort('arnold', 'Surface')

        print '\n--------------------------------- Create Nodes -------------------------------------------\n'

        #Create displacement shader
        if hasDisp:
            network_knode.addShaderInputPort('arnold', 'Displacement')
            involvedNodes = []
            #The name of the diplacement node
            dispShader_name = sg.getElementsByTagName('shaderref')[1].getAttribute('name')
            #The index of the diplacement node, and use the id to get the input shader
            indexInXNodeList = getIndexInList(shader_xmlNodes, 'name', dispShader_name)
            #All co-shaders XML elements connected to the diplacement network
            disp_coShaders = shader_xmlNodes[indexInXNodeList].getElementsByTagName('coshader')

            # Get displacement co-shaders node from xml elments
            for coshader in disp_coShaders:

                if coshader.getAttribute('shader') not in involvedNodes:
                    involvedNodes.append(coshader.getAttribute('shader'))

            # Append all shader nodes connected to co-shader by recursion
            if len(involvedNodes) > 0:
                traverseNodeArray(involvedNodes, shader_xmlNodes)

            for coshader in involvedNodes:
                coshader_name = MAT_NAME_SPACE + '_' + coshader
                disCoShader = ArnoldNode(name=coshader_name, type=shaderNameToType[coshader], parent=shadingGroup.node)

                if disCoShader.name == MAT_NAME_SPACE + '_' + disp_coShaders[0].getAttribute('shader'):
                    # Connect the displacement shader to network's displacement input
                    disCoShader.outputPortByIndex(0).connect(network_knode.inputPortByIndex(1))

            #Remove the diplacement xml elment, to prevent throwing errors when making connection
            for nd in shader_xmlNodes:

                if nd.getAttribute('shaderprogram') == 'displacementShader':
                    shader_xmlNodes.remove(nd)

        # -------------------------------------------------------------------------------------------------------
        #Create surface shader
        surfaceShader_name = sg.getElementsByTagName('shaderref')[0].getAttribute('name')
        surfaceShader = ArnoldNode(name=MAT_NAME_SPACE + '_' + surfaceShader_name,
                                   type=shaderNameToType[surfaceShader_name],
                                   parent=shadingGroup.node)
        #The index of the surface shader node, and use the id to get the input shader
        #Set surface shader's parameters
        indexInXNodeList = getIndexInList(shader_xmlNodes, 'name', surfaceShader_name)
        nodeParams_xmlElement = shader_xmlNodes[indexInXNodeList].getElementsByTagName('parameter')
        surfaceShader.setValue(nodeParams_xmlElement)
        # shader_xmlNodes.pop(indexInXNodeList)
        #Connect the surface shader to network's surface input
        surfaceShader.outputPortByIndex(0).connect(network_knode.inputPortByIndex(0))

        # -------------------------------------------------------------------------------------------------------
        # All co-shaders XML elements connected to the surface shader network
        involvedNodes = []
        surface_coShaders = shader_xmlNodes[indexInXNodeList].getElementsByTagName('coshader')
        for coshader in surface_coShaders:

            if coshader.getAttribute('name') not in involvedNodes:
                involvedNodes.append(coshader.getAttribute('shader'))

        if len(involvedNodes) > 0:
            traverseNodeArray(involvedNodes, shader_xmlNodes)

        for coshader in involvedNodes:
            coshader_name = MAT_NAME_SPACE + '_' + coshader

            # print (BColors.OKBLUE + 'Creating node [%s]' % (coshader_name) + BColors.ENDC)

            shader = ArnoldNode(name=coshader_name, type=shaderNameToType[coshader], parent=shadingGroup.node)
            # Set surface shader's parameters
            indexInXNodeList = getIndexInList(shader_xmlNodes, 'name', coshader)
            nodeParams_xmlElement = shader_xmlNodes[indexInXNodeList].getElementsByTagName(
                'parameter')
            shader.setValue(nodeParams_xmlElement)

        # -------------------------------------------------------------------------------------------------------
        # Connect the new shading group to the main merge node.
        mtlXmerge.addInputPort('i%d' %mtlXmerge.numInputPorts).connect(network_knode.outputPort('out'))

        #Append the new shading group node to a list
        shadingGroup_kNodes.append(shadingGroup)

        # -------------------------------------------------------------------------------------------------------
        #Create material assign node
        assignments = GetAssignments(sg.getElementsByTagName('collection'), gnode)
        materialAssign_name = shadingGroup_name + '_Assignment'
        materialAssign = MaterailAssignNode(name=materialAssign_name, type='MaterialAssign', parent=gnode)
        materialAssign.color = [0.17, (x + 0.2) / (len(shadingGroup_xmlNodes) * 2), 0.29]
        materialAssign.expression = shadingGroup_name
        materialAssign.CEL = assignments
        materialAssign_kNodes.append(materialAssign)

    print '\n--------------------------------- Restore Connections -------------------------------------------\n'

    #Reconnect each other
    for index, shader in enumerate(shader_xmlNodes):
        targetShader_name = MAT_NAME_SPACE + '_' + shader.getAttribute('name')

        if len(targetShader_name.split(':')) > 1:
            targetShader_name = targetShader_name.replace(':', '_')

        targetShader_node = ArnoldNode(targetShader_name)
        coshaders = shader_xmlNodes[index].getElementsByTagName('coshader')

        for co in coshaders:
            sourceShader_name = MAT_NAME_SPACE + '_' + co.getAttribute('shader')

            if len(sourceShader_name.split(':')) > 1:
                sourceShader_name = sourceShader_name.replace(':', '_')

            sourceShader_node = ArnoldNode(sourceShader_name)
            sourceParam_name = co.getAttribute('aovset')[2:]

            if ArnoldNode.param_map.has_key(sourceParam_name):
                sourceParam_name = ArnoldNode.param_map[sourceParam_name]


            targetParam_name = co.getAttribute('name')[2:]

            if ArnoldNode.param_map.has_key(targetParam_name):
                #Two or more params
                #color1(Normal): color1(Normal)
                #color1(blendColor): input1(mix_rgba)
                targetParam_name = ArnoldNode.param_map[targetParam_name]
                if isinstance(targetParam_name, list):
                    for n in targetParam_name:
                        if targetShader_node.node.getParameter('parameters.%s' %ReformatParameter(n, False)):
                            targetParam_name = n

            if targetParam_name[-1] in ['R', 'G', 'B', 'X', 'Y', 'Z']:
                targetParam_name = ReformatParameter(targetParam_name, True)

            if sourceShader_node.node.getParameter('parameters.%s' % ReformatParameter(sourceParam_name, False)) != None:
                sourceParam_name = ReformatParameter(sourceParam_name, False)

            elif ArnoldNode.nodeType_map.has_key(sourceParam_name):
                sourceParam_name = ArnoldNode.nodeType_map[sourceParam_name]

            if targetShader_node.node.getParameter('parameters.%s' %ReformatParameter(targetParam_name, False)) != None:
                targetParam_name = ReformatParameter(targetParam_name, False)

            elif ArnoldNode.nodeType_map.has_key(targetParam_name):
                targetParam_name = ArnoldNode.nodeType_map[targetParam_name]

            print (BColors.OKBLUE + 'Connect [%s.%s]\t >>>>> [%s.%s]'
                   % (sourceShader_node.name, sourceParam_name, targetShader_node.name, targetParam_name) + BColors.ENDC)


            # if sourceShader_node.nodeType == 'place2dTexture':
            #     continue

            if targetParam_name[-1] in ['R', 'G', 'B']:
                targetInput_port = targetShader_node.node.getInputPort(targetParam_name[-1])

                if targetInput_port:
                    sourceShader_node.node.getOutputPort(sourceParam_name).connect(targetInput_port)


            elif sourceShader_node.node.getOutputPort(sourceParam_name):
                targetInput_port = targetShader_node.node.getInputPort(targetParam_name)

                if targetInput_port:
                    sourceShader_node.node.getOutputPort(sourceParam_name).connect(targetInput_port)


            else:
                print(BColors.HEADER + 'Failed to connect from [%s.%s]\t >>>>> [%s.%s]'
             % (sourceShader_node.name, sourceParam_name, targetShader_node.name, targetParam_name) + BColors.ENDC)

    print '\n--------------------------------- dsMatX_Resolve finished -------------------------------------------\n'


    # ------------------------------------------------------------------------------------------------------------
    # Connect merge's input port to matXIn's input port
    # superTools.grp.input --> mainMerge.input
    mtlXmerge.addInputPort('i%d' %mtlXmerge.numInputPorts).connect(gnode.getSendPort('in'))

    # Connect merge's output port to matXIn's output port
    # mainMerge.output --> superTools.grp.output
    mtlXmerge.outputPortByIndex(0).connect(gnode.getReturnPort('out'))

    #Connect Main merge to materialAssign nodes
    if len(materialAssign_kNodes) > 0:
        mtlXmerge.outputPortByIndex(0).connect(materialAssign_kNodes[0].inputPort('input'))

        for i in range(len(materialAssign_kNodes) - 1):
            materialAssign_kNodes[i].outputPort('out').connect(materialAssign_kNodes[i+1].inputPort('input'))

        materialAssign_kNodes[-1].outputPort('out').connect(gnode.getReturnPort('out'))

    #------------------------------------------------------------------------------------------------------------
    #Auto arrange shadingGroup and materialAssign nodes
    DrawingModule.AutoPositionNodes(gnode.getChildren())

    for shadingGroup in shadingGroup_kNodes:
        DrawingModule.AutoPositionNodes(shadingGroup.node.getChildren())


def GetAssignments(collectionTag, gnode):
    '''
    @param collectionTag:   'collection' tag in the xml file
    @param gnode:           superTools.node
    @return:                CEL in the materialAssign, a long string
    '''

    assignments = []

    if len(collectionTag):

        for collection in collectionTag:
            collAdd = collection.getElementsByTagName('collectionadd')

            for coll in collAdd:
                geomAdd = str(coll.getAttribute('geom')).split(', ')
                assignments.extend(geomAdd)

    if len(assignments):
        assignments = ['%s%s' % (gnode.getParameter('location').getValue(0), y) for y in assignments]

    return str(string.join(assignments))


def ReformatParameter(stringParameter, isComponent):
    #Reformat parameters as Katana parameters
    #maya: subsurfaceScale      diffuseRoughness        colorR      scaleX
    #kata: subsurface_scale     diffuse_roughness       color.r     scale.x

    #True: .x, .y., .z
    #False: x_y
    tempSplit = [ a for a in re.split('([A-Z][a-z]*\\d*)', stringParameter) if a ]
    combineString = ''
    for partialString in tempSplit:
        partialString = partialString.lower()
        if combineString == '':
            combineString = partialString
        else:
            combineString += '_' + partialString

    if isComponent:
        attrComponent = combineString.rsplit('_', 1)
        return str('.').join(attrComponent)
    else:
        return combineString


def getIndexInList(sourceList, attribute, valueToFind):
    '''
    Get the index of a value in a element list by the tag given
    @param sourceList:      source element list
    @param attribute:       by the attribute
    @param valueToFind:     the value of the attribute
    @return:
    '''
    indexForAttribute = None
    for obj in sourceList:
        if obj.getAttribute(attribute) == valueToFind:
            indexForAttribute = sourceList.index(obj)

    return indexForAttribute


def RestoreToDefault(gnode):

    children = gnode.getChildren()

    for child in children:
        child.delete()


def traverseNodeArray(involvedNodes, nodes):
    for node in involvedNodes:
        nodeIndex = getIndexInList(nodes, 'name', node)
        currNodeCoshaders = nodes[nodeIndex].getElementsByTagName('coshader')

        if len(currNodeCoshaders) > 0:
            indexInShaderList = getIndexInList(nodes, 'name', node)
            nodeCoShaders = nodes[indexInShaderList].getElementsByTagName('coshader')

            for con in nodeCoShaders:

                if con.getAttribute('shader') not in involvedNodes:
                    involvedNodes.append(con.getAttribute('shader'))


# def prefix_place2dTexture(xml):
#     shaders = xml.getElementsByTagName('shader')
#
#     for i, sd in enumerate(shaders):
#
#         if sd.getAttribute('shaderprogram') == 'place2dTexture':
#             #
#             p2d_name = sd
#             p2d_inputPort = 'i_passthrough'
#
#             relativePath = dict()
#
#             # tex_name = None
#             # tex_outputPort = 'o_output'
#             #
#             # shader_name = None
#             # shader_inputPort = ''
#
#
#             for shader in shaders:
#                 coshaders = shader.getElementsByTagName('coshader')
#
#                 for coshader in coshaders:
#                     co_name = coshader.getAttribute('shader')
#                     if co_name == p2d_name.getAttribute('name'):
#                         tex_name = coshader.parentNode
#                         relativePath[tex_name] = []
#
#                         for sdd in shaders:
#                             co_shaders = sdd.getElementsByTagName('coshader')
#
#                             for co in co_shaders:
#                                 co_name = co.getAttribute('shader')
#
#                                 if co_name == tex_name.getAttribute('name'):
#                                     shader_name = sdd
#                                     shader_inputPort = co.getAttribute('name')
#                                     relativePath[tex_name].append([shader_name, co.getAttribute('name')])
#                                     co.parentNode.removeChild(co)
#             print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
#             print relativePath
#             # # p2d -> shader
#             for tex, value in relativePath.items():
#                 p2d_name.setAttribute('name', tex.getAttribute('name') + '_' + p2d_name.getAttribute('name'))
#                 for i in value:
#                     p2d_element = xml.createElement('coshader')
#                     p2d_element.setAttribute('aovset', 'o_output')
#                     p2d_element.setAttribute('name', i[1])
#                     p2d_element.setAttribute('shader', p2d_name.getAttribute('name'))
#                     i[0].appendChild(p2d_element)
#
#             # shader_element = xml.createElement('coshader')
#             # shader_element.setAttribute('aovset', 'o_output')
#             # shader_element.setAttribute('name', shader_inputPort)
#             # shader_element.setAttribute('shader', p2d_name.getAttribute('name'))
#             # shader_name.appendChild(shader_element)
#             #
#             # #tex -> p2d
#             # p2d_element = xml.createElement('coshader')
#             # p2d_element.setAttribute('aovset', 'o_output')
#             # p2d_element.setAttribute('name', p2d_inputPort)
#             # p2d_element.setAttribute('shader', tex_name.getAttribute('name'))
#             # sd.appendChild(p2d_element)
#
#
#     f = open('/home/sunweiwu/usd/ddd.xml', 'w')
#     xml.writexml(f, addindent = ' ' , newl = '\n' ,encoding = 'utf-8')
#     f.close()
#
#     return xml

def prefix_gainOffset(xml):
    shaders = xml.getElementsByTagName('shader')
    param_map = {
        'colorGain': 'multiply',
        'colorOffset': 'add',
        'alphaGain': 'alpha_multiply',
        'alphaOffset': 'alpha_add'
    }

    for i, sd in enumerate(shaders):
        params = sd.getElementsByTagName('parameter')
        shader_elem = None
        for pa in params:
            if pa.getAttribute('name')[2:] in param_map.keys():
                for sdd in shaders:
                    for co in sdd.getElementsByTagName('coshader'):
                        if co.getAttribute('shader') == sd.getAttribute('name'):
                            #Setting coshader for shader
                            co.setAttribute('shader', sd.getAttribute('name')+'_CC')

                            #Create cc node
                            shader_elem = xml.createElement('shader')
                            shader_elem.setAttribute('name', sd.getAttribute('name')+'_CC')
                            shader_elem.setAttribute('shadertype', 'utility')
                            shader_elem.setAttribute('shaderprogram', 'color_correct')
                            sd.parentNode.appendChild(shader_elem)

                            #Setting coshader for cc node
                            co_cc_elem = xml.createElement('coshader')
                            co_cc_elem.setAttribute('aovset', 'o_output')
                            co_cc_elem.setAttribute('name', 'i_input')
                            co_cc_elem.setAttribute('shader', sd.getAttribute('name'))
                            shader_elem.appendChild(co_cc_elem)

                            break

                #Setting parameters for CC node
                param_cc_elem = xml.createElement('parameter')
                param_cc_elem.setAttribute('name', 'i_' + param_map[pa.getAttribute('name')[2:]])
                param_cc_elem.setAttribute('type', 'color3')
                param_cc_elem.setAttribute('value', pa.getAttribute('value'))
                shader_elem.appendChild(param_cc_elem)


    # f = open('/home/sunweiwu/tttt/ddd.xml', 'w')
    # xml.writexml(f, addindent = ' ' , newl = '\n' ,encoding = 'utf-8')
    # f.close()

    return xml
